package SSF.App.DDoS;

import java.io.*;
import java.util.*;

/** Tempararily used  Destination List tool in DDoS trace back.
 * List destination IP addresses, for each pair of src/dest, it's 
 * counted only once.
 */
public class DestList_3 extends BasicPlayer
{

    // !!! FIX ME !!!!! temparary usage
    //int minIP;
    //int maxIP;

    /** some constants*/
    public static final int SET_QUERY_FAIL = -1;
    public static final int SET_QUERY_SUCC = 0;

    protected HashSet srcIPSet = null;

    /** store information of related ingress traffic*/
    protected Hashtable destTable;

    protected HashSet srcDestSet;

    /** start time */
    protected float startTime=0;

    /** end time */
    protected float endTime=Float.MAX_VALUE;

    private int netFlowTypeCode=-1;

    public DestList_3(String streamID)
    {
	super(streamID);
	destTable = new Hashtable();
	srcDestSet = new HashSet();
	
	//minIP = IP_s.StrToInt("0.0.9.0");
	//maxIP = IP_s.StrToInt("0.0.9.255");
    }

    public void assignSet(HashSet ipSet)
    {
	srcIPSet = ipSet;
    }

    /**
     * This is called to set the parameters for query.
     */
    public int setQuery(float startTime, float endTime)
    {
	this.startTime = startTime;
	this.endTime = endTime;

	if (startTime >= endTime) {
	    System.err.println("Invalid startTime/endTime.");
	    return SET_QUERY_FAIL;
	}
	return SET_QUERY_SUCC;
    }
    
    @Override
	public int receive(int type_code, int source_code, double timestamp, 
		       byte[] bytes, int offset, int length)
    {
	//check type
	if (netFlowTypeCode == -1) {
	    netFlowTypeCode = getRecordTypeCode(IpFlowCollector.STREAM_TYPE);
	}
	if (type_code != netFlowTypeCode) {
	    return 0;             // not a netflow record
	}

	// construct flow
	IpNetFlowWD flow = new IpNetFlowWD();
	flow.readFromBytes(bytes, length);

	// check whether it's within the time period
	if (flow.Last < startTime || flow.First > endTime) {
	    return 0;             // out of the time window.
	}

	//if (flow.srcAddr < minIP || flow.srcAddr > maxIP)
	//    return 0;

	if (srcIPSet != null) {
	    //if there is a srcIPSet, check whether the src IP
	    if (!srcIPSet.contains(new Integer(flow.srcAddr)))
		return 0;
	}
	//update the src destination set.
	Long sdKey = new Long(((long)flow.srcAddr)<<32 | flow.dstAddr);


	if (!srcDestSet.contains(sdKey))
	    srcDestSet.add(sdKey);
	return 0;
    }

    public void collectDest()
    {
	Integer dKey;
	Integer siteNum;

	for (Iterator iter = srcDestSet.iterator();iter.hasNext();) {
	    dKey = new Integer(((Long)iter.next()).intValue());
	    //dKey = new Integer((int)(((long)0xffffffff) & ((Long)iter.next()).longValue()));
	    siteNum = (Integer)destTable.get(dKey);

	    if (siteNum != null)
		destTable.put(dKey, new Integer(siteNum.intValue()+1));
	    else 
		destTable.put(dKey, new Integer(1));
	}
    }

    /**
     * print out the content and a summary of ingress/egress table
     */
    public void printSummary()
    {
	Integer ip;
	System.out.println("--------- destination list  ---------");
	TraceInfo tinfo;
	for (Enumeration enum = destTable.keys();
	     enum.hasMoreElements();) {
	    ip = (Integer)enum.nextElement();	    
	    System.out.println("IP: "+IP_s.IPtoString(ip.intValue())
			       +"\tnumSrcHost:"
			       +((Integer)destTable.get(ip)).intValue());
	}
	System.out.println("-------------------------------------");
    }

    /**
     * Create a hashset from a file.
     * In the file, each line is an IP address.
     */
    public static HashSet createSet(String srcIPFileName)
    {
	LineNumberReader lnr = null;
	HashSet hs = new HashSet();
	try {
	    lnr = new LineNumberReader(new FileReader(srcIPFileName));
	    String str = lnr.readLine();

	    while (str!=null) {
		hs.add(new Integer(IP_s.StrToInt(str)));
		str = lnr.readLine();
	    }
	} catch (Exception any) {
	    System.err.println(any.getMessage() + " "+srcIPFileName);
	    any.printStackTrace();
	    return null;
	}
	return hs;
    }


    public static void main(String[] args)
    {
	float startTime = 0;
	float endTime = Float.MAX_VALUE;
	String srcFileName=null;
	int targetIP=0;
	String streamID=null;

	// get argument
	try {
	    switch (args.length) {
	    case 5:
		if (!args[0].equals("-t"))
		    usageError();
		startTime = Float.valueOf(args[1]).floatValue();
		endTime = Float.valueOf(args[2]).floatValue();
		srcFileName = args[3];
		streamID = args[4];
		break;
	    case 2:
		srcFileName = args[0];
		streamID = args[1];
		break;
	    default:
		usageError();
	    }
	} catch (Exception any) {
	    System.err.println("Error in parameter.");
	    System.err.println(any.getMessage());
	    usageError();
	}

	//for debug
	System.out.println("parameter check passed.");

	//!!!FIX ME
	//System.err.println("This version is not complete yet!!!\n"+
	//	   "Please use DestList_2 for now.");

	DestList_3 destList = new DestList_3(streamID);
      	destList.setQuery(startTime, endTime);

	try {
	    destList.connectRead(srcFileName);
	} catch (streamException se) {
	    System.err.println(se.getMessage());
	    se.printStackTrace();
	    System.exit(-1);
	}
	destList.collectDest();
	destList.printSummary();
    }

    private static void usageError()
    {
	System.err.println("usage: java DestList_3 [-t startTime endTime] <netflow file> <streamID>");
	System.exit(-1);
    }
}
